from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import List, Optional, Union
from enum import Enum, auto

from parse.expr import Expr, ExprVariable
from parse.token import Token


class StmtVisitor(ABC):
    @abstractmethod
    def visit_print_stmt(self, stmt: 'StmtPrint'):
        """"""
    @abstractmethod
    def visit_expr_stmt(self, stmt: 'StmtExpr'):
        """"""
    @abstractmethod
    def visit_var_stmt(self, stmt: 'StmtVar'):
        """"""
    @abstractmethod
    def visit_block_stmt(self, stmt: 'StmtBlock'):
        """"""
    @abstractmethod
    def visit_if_stmt(self, stmt: 'StmtIf'):
        """"""
    @abstractmethod
    def visit_while_stmt(self, stmt: 'StmtWhile'):
        """"""
    @abstractmethod
    def visit_break_stmt(self, stmt: 'StmtBreak'):
        """"""
    @abstractmethod
    def visit_continue_stmt(self):
        """"""
    @abstractmethod
    def visit_class_stmt(self, stmt: 'StmtClass'):
        """"""
    @abstractmethod
    def visit_function_stmt(self, stmt: 'StmtFun'):
        """"""
    @abstractmethod
    def visit_return_stmt(self, stmt: 'StmtReturn'):
        """"""


class Stmt(ABC):
    @abstractmethod
    def accept(self, visitor: 'StmtVisitor'):
        """"""


@dataclass
class StmtExpr(Stmt):
    expr: Expr

    def accept(self, visitor: 'StmtVisitor'):
        return visitor.visit_expr_stmt(self)


class StmtPrint(Stmt):

    def __init__(self, expr: Expr):
        self.expr = expr

    def __repr__(self):
        return f'Print({self.expr})'

    def accept(self, visitor: 'StmtVisitor'):
        return visitor.visit_print_stmt(self)


class StmtVar(Stmt):
    def __init__(self, name: Token, initializer: 'Expr'):
        self.name = name
        self.initializer = initializer
        self.local = True  # True-分配在栈上，False-分配在堆上

    def __repr__(self):
        return f'Var({self.name.lexeme} = {self.initializer})'

    def accept(self, visitor: 'StmtVisitor'):
        return visitor.visit_var_stmt(self)


class StmtBlock(Stmt):
    def __init__(self, stmts: List['Stmt']):
        self.stmts = stmts

    def __repr__(self):
        stmt_str = '\n'.join([str(stmt) for stmt in self.stmts])
        return f'Block({stmt_str})'

    def accept(self, visitor: 'StmtVisitor'):
        return visitor.visit_block_stmt(self)


class StmtIf(Stmt):
    def __init__(self, condition: Expr, then_stmt: 'Stmt', else_stmt: Optional['Stmt']):
        self.condition = condition
        self.then_stmt = then_stmt
        self.else_stmt = else_stmt

    def __repr__(self):
        return f'If(condition={self.condition}, \nthen={self.then_stmt}, \nelse={self.else_stmt})'

    def accept(self, visitor: 'StmtVisitor'):
        return visitor.visit_if_stmt(self)


class StmtWhile(Stmt):
    def __init__(self, condition: 'Expr', body: 'Stmt'):
        self.condition = condition
        self.body = body

    def __repr__(self):
        return f'While(condition={self.condition}, body={self.body})'

    def accept(self, visitor: 'StmtVisitor'):
        return visitor.visit_while_stmt(self)


class FunType(Enum):
    FUN = auto()
    METHOD = auto()


@dataclass
class StmtFun(Stmt):
    name: Token
    params: List['Token']
    body: List['Stmt']
    fun_type: 'FunType'

    def accept(self, visitor: 'StmtVisitor'):
        return visitor.visit_function_stmt(self)


@dataclass
class StmtClass(Stmt):
    name: Token
    superclass: Optional[Token]
    methods: List['StmtFun']

    def accept(self, visitor: 'StmtVisitor'):
        return visitor.visit_class_stmt(self)


@dataclass
class StmtReturn(Stmt):
    expr: Union['Expr', None]

    def accept(self, visitor: 'StmtVisitor'):
        return visitor.visit_return_stmt(self)


class StmtBreak(Stmt):
    def accept(self, visitor: 'StmtVisitor'):
        return visitor.visit_break_stmt(self)


class StmtContinue(Stmt):
    def accept(self, visitor: 'StmtVisitor'):
        return visitor.visit_continue_stmt()
